import numpy as np
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 指定默认字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

def lwlr(testPoint, xArr, yArr, k=1.0):
    """
    局部加权线性回归
    testPoint: 测试点
    xArr: 训练特征
    yArr: 训练标签
    k: 带宽参数
    """
    xMat = np.mat(xArr)
    yMat = np.mat(yArr).T
    m = np.shape(xMat)[0]
    
    # 创建权重矩阵
    weights = np.mat(np.eye(m))
    
    # 计算每个样本点的权重
    for j in range(m):
        diffMat = testPoint - xMat[j,:]
        weights[j,j] = np.exp(diffMat * diffMat.T / (-2.0 * k**2))
    
    # 计算加权最小二乘解
    xTx = xMat.T * (weights * xMat)
    if np.linalg.det(xTx) == 0.0:
        print("矩阵不可逆")
        return
    
    ws = xTx.I * (xMat.T * (weights * yMat))
    return testPoint * ws

def lwlrTest(testArr, xArr, yArr, k=1.0):
    """
    对测试集进行局部加权线性回归预测
    """
    m = np.shape(testArr)[0]
    yHat = np.zeros(m)
    for i in range(m):
        yHat[i] = lwlr(testArr[i], xArr, yArr, k)
    return yHat

# 生成示例数据
def generateData():
    np.random.seed(42)
    x = np.linspace(0, 10, 100)
    # 非线性关系：正弦波加噪声
    y = np.sin(x) + 0.5 * np.random.randn(100)
    return x, y

# 使用示例
x, y = generateData()
xArr = np.column_stack([np.ones(len(x)), x])  # 添加偏置项
yArr = y

# 创建测试点
testX = np.linspace(0, 10, 100)
testArr = np.column_stack([np.ones(len(testX)), testX])

# 不同带宽参数的比较
plt.figure(figsize=(15, 5))

# k = 0.1
plt.subplot(1, 3, 1)
yHat1 = lwlrTest(testArr, xArr, yArr, k=0.1)
plt.scatter(x, y, s=10, alpha=0.5)
plt.plot(testX, yHat1, 'r-', linewidth=2)
plt.title('LWLR with k=0.1 (过拟合)')

# k = 1.0
plt.subplot(1, 3, 2)
yHat2 = lwlrTest(testArr, xArr, yArr, k=1.0)
plt.scatter(x, y, s=10, alpha=0.5)
plt.plot(testX, yHat2, 'r-', linewidth=2)
plt.title('LWLR with k=1.0 (适中)')

# k = 5.0
plt.subplot(1, 3, 3)
yHat3 = lwlrTest(testArr, xArr, yArr, k=5.0)
plt.scatter(x, y, s=10, alpha=0.5)
plt.plot(testX, yHat3, 'r-', linewidth=2)
plt.title('LWLR with k=5.0 (欠拟合)')

plt.tight_layout()
plt.show()